package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"

	"gopkg.in/yaml.v3"
)

// 全局变量存储规则
var (
	rules     []Rule
	rulesLock sync.Mutex
)

// 转发规则
type Rule struct {
	ListenPort int    `yaml:"listen_port"` // 代理监听端口
	TargetAddr string `yaml:"target_addr"` // 后端目标地址（IP:端口）
}

// 配置文件结构
type Config struct {
	Rules []Rule `yaml:"rules"`
}

func loadConfigFromFile(filePath string) (Config, error) {
	var config Config
	data, err := os.ReadFile(filePath)
	if err != nil {
		return config, err
	}
	err = yaml.Unmarshal(data, &config)
	if err != nil {
		fmt.Printf("解析 YAML 失败: %v\n", err)
		return config, err
	}
	return config, nil
}

func loadConfigFromEnv() []Rule {
	var rules []Rule
	for _, env := range os.Environ() {
		if strings.HasPrefix(env, "PROXY_RULES") {
			ruleStr := strings.TrimPrefix(env, "PROXY_RULES=")
			parts := strings.SplitN(ruleStr, ":", 3)
			if len(parts) != 3 {
				fmt.Printf("无效的环境变量规则格式: %s\n", ruleStr)
				continue
			}
			listenPort, err := strconv.Atoi(parts[1])
			if err != nil {
				fmt.Printf("无法解析监听端口: %s\n", parts[1])
				continue
			}
			rule := Rule{
				ListenPort: listenPort,
				TargetAddr: parts[2],
			}
			rules = append(rules, rule)
		}
	}
	return rules
}

func deduplicateRules(rules []Rule) []Rule {
	ruleMap := make(map[string]bool)
	var uniqueRules []Rule
	for _, rule := range rules {
		key := fmt.Sprintf("%d:%s", rule.ListenPort, rule.TargetAddr)
		if _, exists := ruleMap[key]; !exists {
			ruleMap[key] = true
			uniqueRules = append(uniqueRules, rule)
		}
	}
	return uniqueRules
}

func updateRules(newRules []Rule) {
	rulesLock.Lock()
	defer rulesLock.Unlock()
	ruleMap := make(map[int]int)
	for i, rule := range rules {
		ruleMap[rule.ListenPort] = i
	}
	for _, newRule := range newRules {
		if idx, exists := ruleMap[newRule.ListenPort]; exists {
			rules[idx].TargetAddr = newRule.TargetAddr
		} else {
			rules = append(rules, newRule)
		}
	}
}

func getRules() []Rule {
	rulesLock.Lock()
	defer rulesLock.Unlock()
	return rules
}

func httpHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "POST":
		var newRules []Rule
		err := json.NewDecoder(r.Body).Decode(&newRules)
		if err != nil {
			http.Error(w, "无效的请求体", http.StatusBadRequest)
			return
		}
		updateRules(newRules)
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("规则已更新"))
	case "GET":
		currentRules := getRules()
		rulesJSON, err := json.Marshal(currentRules)
		if err != nil {
			http.Error(w, "无法序列化规则", http.StatusInternalServerError)
			return
		}
		w.Header().Set("Content - Type", "application/json")
		w.Write(rulesJSON)
	default:
		http.Error(w, "不支持的方法", http.StatusMethodNotAllowed)
	}
}

func main() {
	// 检查并创建logs目录
	if _, err := os.Stat("./logs"); os.IsNotExist(err) {
		if err := os.MkdirAll("./logs", 0755); err != nil {
			log.Fatalf("无法创建logs目录: %v", err)
		}
	}
	// 打开日志文件
	logFile, err := os.OpenFile("./logs/tcp-proxy.log", os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		log.Fatalf("无法打开日志文件: %v", err)
	}
	defer logFile.Close()

	multiWriter := io.MultiWriter(logFile, os.Stdout)

	// 设置日志输出为MultiWriter
	log.SetOutput(multiWriter)

	config, err := loadConfigFromFile("tcp_proxy.yaml")
	if err != nil {
		fmt.Printf("无法从文件加载配置: %v\n", err)
		return
	}

	envRules := loadConfigFromEnv()
	allRules := append(config.Rules, envRules...)
	uniqueRules := deduplicateRules(allRules)

	updateRules(uniqueRules)

	// 启动HTTP服务器
	go func() {
		http.HandleFunc("/rules", httpHandler)
		log.Fatal(http.ListenAndServe(":8084", nil))
	}()

	var wg sync.WaitGroup
	ctx, cancel := context.WithCancel(context.Background())

	// 为每个规则启动一个监听器
	for _, rule := range uniqueRules {
		wg.Add(1)
		go func(rule Rule) {
			defer wg.Done()
			listenAddr := fmt.Sprintf("0.0.0.0:%v", rule.ListenPort)
			// 启动 TCP 监听
			listener, err := net.Listen("tcp", listenAddr)
			if err != nil {
				log.Printf("监听端口 %d 失败: %v", rule.ListenPort, err)
				return
			}
			defer listener.Close()

			log.Printf("TCP 反向代理启动: 本地:%d -> 后端:%s", rule.ListenPort, rule.TargetAddr)

			for {
				select {
				case <-ctx.Done():
					return
				default:
					clientConn, err := listener.Accept()
					if err != nil {
						if ctx.Err() != nil {
							return
						}
						log.Printf("接收连接失败: %v", err)
						continue
					}
					go handleTCPConnection(ctx, clientConn, rule.TargetAddr)
				}
			}
		}(rule)
	}

	// 捕获系统信号，实现优雅关闭
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-sigCh
		log.Println("收到终止信号:", sig)
		cancel()
		os.Exit(1)
	}()

	wg.Wait()
}

// 处理单个 TCP 连接，实现双向转发
func handleTCPConnection(ctx context.Context, clientConn net.Conn, targetAddr string) {
	defer clientConn.Close()
	log.Printf("新连接: %s -> 后端 %s", clientConn.RemoteAddr(), targetAddr)

	// 连接后端服务
	dialer := &net.Dialer{}
	backendConn, err := dialer.DialContext(ctx, "tcp", targetAddr)
	if err != nil {
		log.Printf("连接后端 %s 失败: %v", targetAddr, err)
		return
	}
	defer backendConn.Close()

	var wg sync.WaitGroup
	wg.Add(2)

	// 客户端 -> 后端
	go func() {
		defer wg.Done()
		_, err := io.Copy(backendConn, clientConn)
		if err != nil && err != io.EOF {
			log.Printf("客户端 -> 后端 转发错误: %v", err)
			backendConn.Close()
		}
	}()

	// 后端 -> 客户端
	go func() {
		defer wg.Done()
		_, err := io.Copy(clientConn, backendConn)
		if err != nil && err != io.EOF {
			log.Printf("后端 -> 客户端 转发错误: %v", err)
			clientConn.Close()
		}
	}()

	// 等待双向转发完成（任一方向关闭则退出）
	go func() {
		<-ctx.Done()
		clientConn.Close()
		backendConn.Close()
	}()

	wg.Wait()
	log.Printf("连接关闭: %s -> 后端 %s", clientConn.RemoteAddr(), targetAddr)
}
